import random
import time
from collections import Counter
import numpy as np
import matplotlib.pyplot as plt

# ========== STEP 1: GIFT Encryption and Plaintext-Ciphertext Pair Generation ==========

def gift_encrypt(plaintext, key, rounds):
    """Simplified GIFT encryption (substitute with actual implementation)."""
    # Implement S-box, key addition, and permutation layers
    ciphertext = plaintext ^ key  # Example operation (Replace with real GIFT rounds)
    return ciphertext

# Generate plaintext-ciphertext pairs
def generate_plaintext_ciphertext_pairs(n_pairs, key, rounds):
    pairs = []
    for _ in range(n_pairs):
        plaintext = random.randint(0, 0xFFFF)
        ciphertext = gift_encrypt(plaintext, key, rounds)
        pairs.append((plaintext, ciphertext))
    return pairs

# ========== STEP 2: Differential Distribution Calculation ==========

def calculate_differences(pairs):
    """Calculate plaintext and ciphertext differences."""
    plaintext_diffs = []
    ciphertext_diffs = []
    for (p1, c1), (p2, c2) in zip(pairs[:-1], pairs[1:]):
        plaintext_diffs.append(p1 ^ p2)
        ciphertext_diffs.append(c1 ^ c2)
    return Counter(plaintext_diffs), Counter(ciphertext_diffs)

# Generate Diffusion Distribution Table (DDT) for the S-box
def generate_ddt(sbox):
    size = len(sbox)
    ddt = np.zeros((size, size), dtype=int)
    for x1 in range(size):
        for x2 in range(size):
            diff_in = x1 ^ x2
            diff_out = sbox[x1] ^ sbox[x2]
            ddt[diff_in][diff_out] += 1
    return ddt

# ========== STEP 3: Traditional Differential Attack ==========

def traditional_attack(pairs, ddt, target_diff):
    """Traditional differential attack simulation."""
    potential_keys = []
    for candidate_key in range(0xFFFF):  # Example: Brute force all possible keys
        matches = 0
        for plaintext, ciphertext in pairs:
            decrypted = ciphertext ^ candidate_key  # Example decryption step
            if decrypted == plaintext ^ target_diff:
                matches += 1
        if matches > len(pairs) * 0.7:  # Example threshold
            potential_keys.append(candidate_key)
    return potential_keys

# ========== STEP 4: MILP Optimization and Improved Attack ==========

def milp_optimized_attack(pairs, ddt, target_diff):
    """Improved attack using MILP-optimized paths."""
    # Placeholder for MILP implementation (e.g., use Gurobi or similar tools)
    potential_keys = []
    for candidate_key in range(0xFFFF):  # Example: Brute force with MILP optimization
        matches = 0
        for plaintext, ciphertext in pairs:
            decrypted = ciphertext ^ candidate_key
            if decrypted == plaintext ^ target_diff:
                matches += 1
        if matches > len(pairs) * 0.8:  # Example improved threshold
            potential_keys.append(candidate_key)
    return potential_keys

# ========== STEP 5: Path Probability Calculation ==========

def calculate_path_probability(path, ddt):
    """Calculate the probability of a differential path."""
    prob = 1.0
    for (delta_in, delta_out) in path:
        if delta_in < len(ddt) and delta_out < len(ddt[0]):
            prob *= ddt[delta_in][delta_out] / sum(ddt[delta_in])
        else:
            prob = 0  # Invalid path, set probability to 0
            break
    return prob

# ========== Visualization and Saving Results ==========

def visualize_results(success_rates, probabilities, times):
    """Visualize success rates, path probabilities, and efficiency."""
    labels = ["Traditional Attack", "Improved Attack"]
    
    # Success Rates
    plt.bar(labels, success_rates, color=["blue", "green"])
    plt.title("Success Rate Comparison")
    plt.ylabel("Success Rate (%)")
    plt.savefig('/home/zcw/gift/success_rate_comparison.png')  # Save the figure
    plt.close()
    
    # Path Probabilities
    plt.bar(labels, probabilities, color=["orange", "purple"])
    plt.title("Path Probability Comparison")
    plt.ylabel("Probability")
    plt.savefig('/home/zcw/gift/path_probability_comparison.png')  # Save the figure
    plt.close()
    
    # Efficiency (Time)
    plt.bar(labels, times, color=["red", "cyan"])
    plt.title("Efficiency Comparison (Time)")
    plt.ylabel("Average Time (s)")
    plt.savefig('/home/zcw/gift/efficiency_comparison.png')  # Save the figure
    plt.close()

# ========== Main Experiment ==========

if __name__ == "__main__":
    # Parameters
    n_pairs = 1000
    rounds = 5
    key = 0x1234
    target_diff = 0x5678  # Example difference
    sbox = [0x6, 0x4, 0xC, 0x5, 0x0, 0x7, 0x2, 0xE, 0x1, 0xF, 0x3, 0xD, 0x8, 0xA, 0x9, 0xB]  # Example S-box

    # Step 1: Generate plaintext-ciphertext pairs
    pairs = generate_plaintext_ciphertext_pairs(n_pairs, key, rounds)
    
    # Save the plaintext-ciphertext pairs to a text file
    with open('/home/zcw/gift/plaintext_ciphertext_pairs.txt', 'w') as f:
        for plaintext, ciphertext in pairs:
            f.write(f"{hex(plaintext)} {hex(ciphertext)}\n")
    
    # Step 2: Calculate differences
    pt_diffs, ct_diffs = calculate_differences(pairs)
    
    # Save the plaintext and ciphertext differences to text files
    with open('/home/zcw/gift/plaintext_diffs.txt', 'w') as f:
        for diff, count in pt_diffs.items():
            f.write(f"{hex(diff)} {count}\n")
    
    with open('/home/zcw/gift/ciphertext_diffs.txt', 'w') as f:
        for diff, count in ct_diffs.items():
            f.write(f"{hex(diff)} {count}\n")
    
    # Step 3: Generate Diffusion Distribution Table (DDT)
    ddt = generate_ddt(sbox)
    
    # Step 4: Traditional attack
    start_time = time.time()
    traditional_keys = traditional_attack(pairs, ddt, target_diff)
    traditional_time = time.time() - start_time
    
    # Step 5: MILP attack (simulated)
    start_time = time.time()
    improved_keys = milp_optimized_attack(pairs, ddt, target_diff)
    improved_time = time.time() - start_time
    
    # Step 6: Calculate path probability (for comparison)
    traditional_path = [(0x1, 0x4), (0x4, 0x6)]
    improved_path = [(0x1, 0x5), (0x5, 0x7)]
    prob_traditional = calculate_path_probability(traditional_path, ddt)
    prob_improved = calculate_path_probability(improved_path, ddt)
    
    # Step 7: Visualization
    visualize_results(
        [len(traditional_keys) / 0xFFFF * 100, len(improved_keys) / 0xFFFF * 100],
        [prob_traditional, prob_improved],
        [traditional_time, improved_time]
    )

